/*
 *  Copyright (C) 2010 Martin Haulrich <mwh.isv@cbs.dk>
 *
 *  This file is part of the MatrixParser package.
 *
 *  The MatrixParser program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.osdtsystem.matrixparser.main;


import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
import java.io.Writer;
import java.util.ArrayList;
import java.util.List;
import org.osdtsystem.matrixparser.data.CONLLIterator;
import org.osdtsystem.matrixparser.data.CONLLSentence;
import org.osdtsystem.matrixparser.data.GraphIterator;
import org.osdtsystem.matrixparser.featureextraction.AbstractFeatureExtractor;
import org.osdtsystem.matrixparser.parsers.output.DependencyTree;
import org.osdtsystem.matrixparser.featureextraction.FeatureExtractor;
import org.osdtsystem.matrixparser.featureextraction.FeatureExtractorRegistry;
import org.osdtsystem.matrixparser.featureextraction.UnionExtractor;
import org.osdtsystem.matrixparser.learners.Learner;
import org.osdtsystem.matrixparser.learners.LinearScorer;
import org.osdtsystem.matrixparser.learners.MIRALearner;
import org.osdtsystem.matrixparser.parsers.ParsingModel;
import org.osdtsystem.matrixparser.learners.Scorer;
import org.osdtsystem.matrixparser.learners.PerceptronLearner;
import org.osdtsystem.matrixparser.logging.DebugData;
import org.osdtsystem.matrixparser.logging.Log;
import org.osdtsystem.matrixparser.parsers.IncrementalLabMSTParser;
import org.osdtsystem.matrixparser.parsers.LabelledMSTParser;
import org.osdtsystem.matrixparser.parsers.Parser;
import org.osdtsystem.matrixparser.parsers.Trainer;
import org.osdtsystem.matrixparser.parsers.loss.LabelledLoss;
import org.osdtsystem.matrixparser.parsers.loss.Loss;

/**
 * Main class for training and running MST and IRP parsers.
 * @author Martin Haulrich and Matthias Buch-Kromann
 */
public class IrpMain {
    public static void main(String[] args) throws Exception {
        Options.parseArgs(args);
        Log.chandler.setLevel(Options.levelOption("log.level"));
        Log.config("Running IrpMain");
        run();
    }

    static void run() throws Exception {
        // Create models, extractors, scorers, parsers, and trainers
        ParsingModel model = makeModel();
        UnionExtractor extractor = makeFeatureExtractor(model);
        Scorer scorer = makeScorer();
        Parser trainParser = makeTrainParser(scorer, model, extractor);
        List<Parser> evalParsers = makeEvalParsers(scorer, model, extractor);
        Trainer trainer = makeTrainer(trainParser, model);

        // Create treebank iterators
        GraphIterator trainIterator = makeGraphIterator(Options.option("train.file"));
        GraphIterator devIterator = makeGraphIterator(Options.option("train.file.dev"));
        List<GraphIterator> evalIterators = makeEvalGraphIterators();

        // Grow alphabet
        model.growAlphabet(extractor, trainIterator);

        // Train parser
        if (trainer != null) {
            trainer.train(trainIterator, devIterator);
            writeModel(model);
        }

        // Evaluate parsers
        runParsers(model, evalParsers, evalIterators);
    }

    static ParsingModel makeModel() {
        File modelFile = Options.fileOption("model.ifile");
        ParsingModel model = new ParsingModel();
        if (modelFile != null && modelFile.canRead()) {
            try {
                Log.config("Loading parsing model from file \"", modelFile.getName(), "\"");
                model = ParsingModel.load(modelFile);
                model.fixWeights();
            } catch (IOException ex) {
                throw new RuntimeException("Cannot open model file: " + modelFile, ex);
            } catch (ClassNotFoundException ex) {
                throw new RuntimeException("Cannot decode model file: " + modelFile, ex);
            }
        }
        return DebugData.model = model;
    }

    private static UnionExtractor makeFeatureExtractor(ParsingModel model) {
        // Create extractor
        FeatureExtractor.deadGenes.or(Options.bitsetOption("genes"));
        UnionExtractor extractor = new FeatureExtractorRegistry(model,
                Options.option("extractors")).unionExtractor();

        // Log gene use
        Log.config("Dead genes: ", Options.bitString(FeatureExtractor.deadGenes),
                " ", FeatureExtractor.deadGenes.toString(),
                " inverted: ", AbstractFeatureExtractor.invertBitString(
                    FeatureExtractor.deadGenes).toString());
        Log.info(extractor.geneUse());

        // Return
        return extractor;
    }

    static Scorer makeScorer() {
        String scorerType = Options.option("scorer.type");
        if (scorerType != null && ! scorerType.equals("linear")) {
            Log.warning("Unknown scorer type \"", scorerType,
                "\": using linear scorer instead");
        }
        Log.info("Using linear scorer");
        return new LinearScorer();
    }

    static Parser makeParser(Scorer scorer, ParsingModel model, UnionExtractor extractor,
            String name) {
        if (name.equals("mst")) {
            return new LabelledMSTParser(name, scorer, extractor, model);
        } else if (name.equals("irp")) {
            return new IncrementalLabMSTParser(name, scorer, model, extractor);
        } else {
            Log.warning("Unknown parser type: ", name);
            return null;
        }
    }

    static Parser makeTrainParser(Scorer scorer, ParsingModel model, UnionExtractor extractor) {
        return makeParser(scorer, model, extractor, Options.option("train.parser"));
    }

    static List<Parser> makeEvalParsers(Scorer scorer, ParsingModel model,
            UnionExtractor extractor) {
        List<Parser> parsers = new ArrayList<Parser>();
        for (String name : Options.option("eval.parsers").split(":")) {
            Parser parser = makeParser(scorer, model, extractor, name);
            if (parser != null)
                parsers.add(parser);
        }
        return parsers;
    }

    private static Trainer makeTrainer(Parser trainParser, ParsingModel model) {
        Loss loss = makeLoss();
        Learner learner = makeLearner(model);
        return new Trainer(learner, trainParser, loss);
    }

    static Loss makeLoss() {
        String lossName = Options.option("loss.type");
        Double wrongHeadLoss = Options.doubleOption("loss.wrongHeadLoss");
        Double wrongLabelLoss = Options.doubleOption("loss.wrongLabelLoss");
        if (lossName != null && ! lossName.equals("labelled")) {
            Log.warning("Unknown loss type \"", lossName,
                    "\": using labelled scorer instead");
        }
        Log.config("Using labelled loss with wrongHeadLoss=",
                String.format("%6g", wrongHeadLoss),
                " and wrongLabelLoss=", String.format("%6g", wrongLabelLoss));
        return new LabelledLoss(wrongHeadLoss, wrongLabelLoss);
    }

    static Learner makeLearner(ParsingModel model) {
        String learnerType = Options.option("learner.type");
        Learner learner = new MIRALearner(model);
        if (learnerType != null) {
            if (learnerType.equals("perceptron")) {
                learner = new PerceptronLearner(0);
            } else {
                Log.warning("Unknown learner type: ", learnerType);
            }
        }
        return learner;
    }

    private static GraphIterator makeGraphIterator(String fileName) {
        if (fileName.endsWith(".conll") || Options.option("fileformat").equals("conll")) {
            try {
                return new CONLLIterator(fileName);
            } catch (FileNotFoundException ex) {
                Log.warning("File not found: ", fileName);
            } catch (UnsupportedEncodingException ex) {
                Log.warning("Cannot decode file: ", fileName);
            }
        }
        return null;
    }

    private static List<GraphIterator> makeEvalGraphIterators() {
        List<GraphIterator> evalIterators = new ArrayList<GraphIterator>();
        for (String name : Options.option("eval.files").split(":")) {
            GraphIterator iterator = makeGraphIterator(name);
            if (iterator != null)
                evalIterators.add(iterator);
        }
        return evalIterators;
    }

    private static void runParsers(ParsingModel model, List<Parser> evalParsers,
            List<GraphIterator> evalIterators) {
        Log.info("Parsing...");

        for (Parser parser : evalParsers) {
            for (GraphIterator evalIterator : evalIterators) {
                try {
                    // Log actions
                    Log.info("Parsing corpus ", evalIterator.name(), " with parser ",
                            parser.name());
                    Log.pushTimeStamp();
                    
                    // Reset iterator
                    evalIterator.reset();
                    Writer writerCONLL = makeOutputWriterCONLL(evalIterator, parser);
                    Writer writerDTAG = makeOutputWriterDTAG(evalIterator, parser);
                    
                    // Parse corpus
                    int i = 0;
                    while (evalIterator.hasNext()) {
                        // Read sentence
                        CONLLSentence sentence = evalIterator.next();
                        if (i++ % 100 == 0) Log.infoW(".");

                        // Parse sentence and write it to conll file
                        DependencyTree parse = parser.parse(sentence, model.weights());
                        if (writerCONLL != null) parse.writeAsCONLL(writerCONLL);
                        if (writerDTAG != null) parse.writeAsDTAG(writerDTAG);
                    }
                    // Log again
                    Log.info("   finished parsing: [", Log.popTimeStamp(), " seconds,  ",
                            Log.heapUsage(), " MB heap]");

                    // Close writers
                    if (writerCONLL != null) writerCONLL.close();
                    if (writerDTAG != null) writerDTAG.close();
                } catch (IOException ex) {
                    Log.warning("Input-outpur error encountered during parsing: ", 
                            ex.toString());
                }
            }
        }
    }



    private static void writeModel(ParsingModel model) {
        File ofile = new File(Options.option("model.ofile"));
        if (ofile != null) {
            try {
                model.save(ofile);
            } catch (IOException ex) {
                Log.severe("Cannot open model file ", ofile.getName(), " for writing");
            }
        }
    }


    private static Writer makeOutputWriterCONLL(GraphIterator evalIterator, Parser parser) {
        throw new UnsupportedOperationException("Not yet implemented");
                Writer writer = new PrintWriter(outputFile, "UTF-8");
                Writer writerDTAG = new PrintWriter(outputFile + ".dtag", "UTF-8");
    }

    private static Writer makeOutputWriterDTAG(GraphIterator evalIterator, Parser parser) {
        throw new UnsupportedOperationException("Not yet implemented");
    }



    /*
     private static void readComplements(Parser trainParser, Parser[] evalParsers) {
        throw new UnsupportedOperationException("Not yet implemented");

                // Read complement file and notify parsers
        String complementFile = Options.option("complements.file");
        if (complementFile != null) {
            int lastComplementId = AbstractParser.readComplements(complementFile, model.labelHandler());
            mstParser.setLastComplementID(lastComplementId);
            incParser.setLastComplementID(lastComplementId);
        }
    }
     */



}

